



import stock_loader

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
from tqdm import tqdm

class MultiStockDataset(Dataset):
    """处理多支股票数据的数据集类"""
    def __init__(self, data_list, seq_length, pred_length=1):
        """
        data_list: 包含多支股票数据的列表，每支股票是一个多维numpy数组[天数，特征矢量]，一个特征矢量索引依次对应[ 开盘价，最高价，最低价，收盘价，成交量，成交额 ]。
        seq_length: 输入序列长度，
        pred_length: 预测序列长度
        """
        self.data_list = data_list
        self.seq_length = seq_length
        self.pred_length = pred_length
        self.stock_num = len(data_list) # 股票的数量
        self.sequences = self._create_sequences()
        
    def _create_sequences(self): # 数据预处理：创建数据集序列。

        sequences = []   # 列表元素是元组(输入序列，预测输出序列，股票索引)

        # enumerate() 函数用于将一个可遍历的数据对象（如列表、元组或字符串）组合为一个索引序列，同时列出数据和数据下标。它通常用于 for 循环中，以便在遍历序列时获取元素的索引和值。
        for stock_idx, stock_data in enumerate(self.data_list): # 这里获取股票索引与数据

            # 股票数据应为 [天数, 特征数] 的numpy数组
            total_days = len(stock_data)  # 这里获取一支股票的长度。
            # 确保有足够的数据进行序列生成和预测
            for i in range(total_days - self.seq_length - self.pred_length + 1):
                # 输入序列：前seq_length天的所有特征
                x = stock_data[i:i+self.seq_length]
                # 目标序列：接下来pred_length天的最高价、最低价和成交量
                y = stock_data[i+self.seq_length:i+self.seq_length+self.pred_length, : ]  
                sequences.append((x, y, stock_idx))  # 同时保存股票索引
        return sequences
    
    #  PyTorch 的 Dataset 类是 PyTorch 处理自定义数据集的标准方式，需要实现__len__和__getitem__方法。

    '''
        Dataset 使用示例：
        dataset = MultiStockDataset(idata_list=[...], seq_length=300)
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)

        # 迭代获取批量样本
        for batch in dataloader:
            inputs, labels = batch 
            # inputs形状: [32, 输入序列], labels形状: [32, 输出序列]
            # 32 是批量大小。
    '''


    def __len__(self):
        return len(self.sequences)
    
    def __getitem__(self, idx): # 根据索引获取单个样本。
        x, y, stock_idx = self.sequences[idx]
        return torch.FloatTensor(x), torch.FloatTensor(y), stock_idx




class PositionalEncoding(nn.Module):
    """位置编码器，为Transformer添加位置信息"""
    def __init__(self, d_model, max_len=5000):
        """
            d_model：模型的隐藏维度，即一个位置矢量的维度。一般等于序列样本的维度, 如果不等于，在Transform就要将样本做线性变换，转换到与d_model一样的维度。
            max_len：位置矢量支持的最长位置。
        """
        super(PositionalEncoding, self).__init__()

        # 创建位置编码矩阵
        '''
            对于位置pos和维度i，位置编码值计算如下：
            PE(pos,2i) = sin( pos / 10000^(2i/d_model) )
            PE(pos,2i+1) = cos( pos / 10000^(2i/d_model) )

            pos：序列中的位置（范围：0 到max_len-1）。
            i：编码维度（范围：0 到d_model-1）。
                   
        '''
        pe = torch.zeros(max_len, d_model)

        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        # torch.unsqueeze() 是一个用于增加张量维度的核心函数，它可以在指定位置插入一个新的维度（大小为 1），不改变原始数据，仅调整形状。常用于调整张量形状以满足运算或模型输入的要求。新维度插入后，原维度依次后移。 
        # 这里是将形状[max_len]的张量变成[max_len,1]  


        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        '''
        详解：
        1、torch.arange(0, d_model, 2) 生成从 0 开始、步长为 2 的序列，对应偶数维度索引 [0, 2, 4, ...]
        2、-np.log(10000.0) / d_model 是一个常量缩放因子，用于控制不同维度的频率。这里其实是得到 -np.log(10000.0)*2i/d_model的序列。
        3、torch.exp(...) 将上述结果转换为指数形式。

        好处：log(10000.0)约等于9.21，数值稳定性好。将大数转成小数运算，数值比较友好。
        '''

        #填充偶数维度（使用正弦函数）
        pe[:, 0::2] = torch.sin(position * div_term)
        #填充奇数维度（使用余弦函数）
        pe[:, 1::2] = torch.cos(position * div_term)
        '''
        1::2 的含义:
        这是 Python 切片的标准语法：start:end:step（end省略时表示到末尾）。
        start=1：从索引1开始（第 2 列，因为索引从 0 开始）。
        step=2：步长为 2，即每隔 1 列选取一次（跳过 1 个元素）
        '''

        # 将位置编码注册为缓冲区，不会被视为模型参数，这个位置编码一次性计算好，放缓冲区。这个序列要保证不大于5000.
        #缓冲区：类似于模型参数（nn.Parameter），但不参与梯度计算和优化。
        #典型用途：存储不需要学习但需要随模型一起保存 / 加载的状态，如 BatchNorm 的统计量、位置编码等。
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        '''
         输入张量 x：形状为 [batch_size, seq_length, d_model]，表示一批序列数据。
         输出张量：形状与输入相同 [batch_size, seq_length, d_model]，但每个元素已添加位置编码信息。
        '''
        # 解析：
        # self.pe：预计算的位置编码矩阵，形状为 [max_len, d_model]
        # x.size(1)：当前输入序列的长度 seq_length。
        # 这里截取位置编码矩阵的前seq_length行，得到形状为 [seq_length, d_model] 的子矩阵。
        # [seq_length, d_model] 的子矩阵前面要增加一个维度，才能与输入相加。
        '''
        加法方式添加位置信息的好处：
        1、维度保持为d_model，参数数量和计算复杂度不变，模型效率更高。
        2、原始 Transformer 论文Vaswani 等人在《Attention Is All You Need》中明确选择加法，并通过实验验证其有效性。大量后续研究（如 BERT、GPT 等）沿用了加法设计，进一步证明其在各种任务中的有效性。拼接方式在实践中很少被采用，因为它通常导致性能下降。
        3、加法更利于捕捉相对位置：位置编码的核心目标是让模型学习序列中元素的相对位置关系。 加法能将序列样本中特征矢量的区别，增加位置的区别。  
        '''
        return x + self.pe[ :x.size(1) ,  : ].unsqueeze(0)



class TransformerModel(nn.Module):
    """完整的Transformer模型，用于股票数据预测"""
    def __init__(self, input_dim, d_model, nhead, num_layers, output_dim, pred_length, dropout=0.1):
        """
        input_dim：   输入特征维度（如股票数据的开盘价、收盘价等特征数量）。
        d_model：     Transformer 隐藏维度，即每个位置的表示维度。
        nhead：多头注意力的头数，并行处理多个子空间的信息。
        num_layers：Transformer 编码器层数，增加深度可捕获更复杂的模式。
        output_dim：预测的目标特征维度（如预测最高价、最低价、成交量）。
        pred_length：预测未来的时间步数（如预测未来 3 天的股价）。
        """

        super(TransformerModel, self).__init__()
        self.input_dim = input_dim
        self.d_model = d_model
        self.pred_length = pred_length
        
        # 输入投影层：将输入特征映射到d_model维度(位置编码的维度)，方便添加位置信息。
        self.input_proj = nn.Linear(input_dim, d_model)
        
        # 位置编码器，输入位置信息。
        self.pos_encoder = PositionalEncoding(d_model)
        
        # Transformer编码器层
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,   #将输入分为多个头，并行计算注意力，捕获不同子空间的依赖关系。
            dim_feedforward=4*d_model,  # 通常设置为d_model的4倍
            dropout=dropout,
            batch_first=True  # PyTorch 1.9+支持batch_first
        )
        '''
        输入张量通常有三个维度：[batch_size, sequence_length, feature_dim] 或 [sequence_length, batch_size, feature_dim]。
        batch_first=False（默认）：PyTorch 传统约定，输入形状为 [seq_len, batch_size, d_model]。
            seq_len：序列长度（如句子中的词数）。
            batch_size：批次大小（并行处理的样本数）。
            d_model：特征维度（如词向量维度）。
        batch_first=True：更符合直觉的约定（类似 TensorFlow/Keras），输入形状为 [batch_size, seq_len, d_model]。
        '''

        '''
        nn.TransformerEncoderLayer 其他参数详解：

        nhead 对应 Transformer 中的 “多头自注意力（Multi-Head Attention）” 机制，其核心是将原始注意力计算分解到多个低维子空间中并行执行，再将结果拼接融合，以捕获更丰富的依赖关系。
            1、设输入序列的特征维度为 d_model（如 512），则每个头的特征维度为 d_k = d_v = d_model / nhead（必须整除，例如 d_model=512、nhead=8 时，每个头的维度为 64）。
            2、对输入矩阵 X（形状 [batch_size, seq_len, d_model]），通过线性变换生成查询（Q）、键（K）、值（V）矩阵：Q=X⋅W_Q ,K=X⋅W_K ,V=X⋅W_V 。 W_Q, W_K, W_V 是可学习参数，形状均为 [d_model, d_model]。
            3、将 Q、K、V 按头拆分（每个头维度 d_k），得到 nhead 组子矩阵(头)，每个子矩阵Q_i,K_i,V_i 形状为 [batch_size, seq_len, d_model / nhead]。
            4、每个头独立计算自注意力：softmax( (Qi*Ki^T)/sqrt(d_k) )*Vi 
            5、最后拼接所有头的结果，再通过线性变换 W_O 映射回 d_model 维度。

            作用
            捕获多尺度依赖：不同头可关注序列中不同位置、不同语义的依赖关系（例如有的头关注局部短语，有的头关注长距离语法）。
            增加模型容量：通过子空间拆分，在相同计算复杂度下（与单头注意力相比），多头机制能建模更复杂的特征交互。

        dim_feedforward（前馈网络的隐藏层维度）.
            Transformer 编码器层的第二个核心模块是前馈神经网络（Feed-Forward Network, FFN），dim_feedforward 定义了该网络隐藏层的维度（本质就是个全连接层）。其结构为： FFN(X) = W2 * ReLu( W1*X + b1 ) + b2

            输入 X 形状为 [batch_size, seq_len, d_model]，先通过线性层 W_1 映射到 dim_feedforward 维度（升维）。
            经过 ReLU 激活函数引入非线性。
            再通过线性层 W_2 映射回 d_model 维度（降维）。

            为什么通常设为 4*d_model？
            这是 Transformer 原论文（Vaswani et al., 2017）的经验设置。
            数学上，升维到 4*d_model 可以为非线性变换提供更大的 “特征空间”，让网络有足够能力学习输入特征的复杂非线性组合（例如高阶特征交互），再通过降维压缩为更有效的表示。
            若维度太小，可能限制非线性拟合能力；太大则会增加计算成本，且可能引入冗余。

        dropout（ dropout 概率）：
            dropout 是一种正则化技术，通过在训练时随机将部分神经元的输出置为 0（概率为 dropout），减少模型对局部特征的过度依赖，防止过拟合。
            在 Transformer 编码器层中，dropout 通常应用于 3 个位置：
                多头注意力输出后（残差连接前）；
                前馈网络输出后（残差连接前）；
                部分实现中会对注意力权重矩阵（QK^T）也应用 dropout。
            设某层输出为 Y，dropout 后结果为 Y'，则Y'的元素中，有p的概率为0，其他的元素缩放成 Yi' = Yi/(1-p), 这时为了保持输出的期望不变。

            作用
            防止过拟合：通过随机“失活”部分神经元，强制模型学习更鲁棒的特征（不依赖特定神经元的输出）。
            增强泛化能力：训练时模型相当于在多个“子网络”间切换，测试时集成所有子网络的效果。
        '''

        
        # Transformer编码器
        # nn.TransformerEncoder 是 PyTorch 中实现 Transformer 编码器堆叠结构的模块，它将多个相同的编码器层（如你之前定义的 encoder_layer）按顺序堆叠，形成深度网络，从而捕获序列中的多层次抽象表示。
        # 层与层之间通过残差连接和层归一化连接。
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        

        # 输出解码器：将Transformer输出映射到预测维度
        self.decoder = nn.Sequential(
            nn.Linear(d_model, 128),
            nn.ReLU(), 
            nn.Dropout(dropout),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, output_dim * pred_length)  # 输出未来pred_length天的预测
        )
        '''
        定义了一个多层感知机（MLP）作为输出解码器，将 Transformer 编码器的输出映射到最终的预测结果。
            Transformer 编码器的输出是高维特征表示（形状为[batch_size, seq_len, d_model]）解码器的作用就是：
            降维：将d_model维度的特征映射到更低维度。
            非线性变换：通过激活函数引入非线性，增强模型表达能力。
            多步预测：输出未来pred_length个时间步的预测值。

            第一层：
            nn.Linear(d_model, 128)：将 Transformer 输出的d_model维特征映射到 128 维。
            nn.ReLU()：引入非线性变换，允许模型学习复杂模式。
            nn.Dropout(dropout)：随机丢弃部分神经元，防止过拟合。

            第二层：
            nn.Linear(128, 64)：进一步降维，压缩特征表示。
            同样使用 ReLU 和 Dropout。

            输出层：
            nn.Linear(64, output_dim * pred_length)：将特征映射到最终输出维度。
            output_dim：每个时间步的预测变量数（如多变量预测中的变量个数）。
            pred_length：预测的未来时间步数。

        假设输入到解码器的张量形状为[batch_size, seq_len, d_model]，解码器的输出形状为：[batch_size, seq_len, output_dim * pred_length]

        从下面forword 函数可以知道，编码器取最后一个时间步的输出作为上下文向量，因此，最终的输出形状是：[batch_size, 1, output_dim * pred_length]
        '''

        
    def forward(self, src):
        # src: [batch_size, seq_length, input_dim]
        batch_size, seq_length, _ = src.shape
        
        # 输入投影，乘以sqrt(d_model)是 Transformer 的常见做法，用于缩放输入，防止梯度消失。d_model是一个常数，所以，np.sqrt(self.d_model)也是个常量。
        src = self.input_proj(src) * np.sqrt(self.d_model)
        
        # 添加位置编码
        src = self.pos_encoder(src)
        
        # 通过Transformer编码器
        transformer_output = self.transformer_encoder(src)
        
        # 取最后一个时间步的输出作为上下文向量
        context = transformer_output[:, -1, :]  # [batch_size, d_model]
        
        '''
        仅使用最后一个时间步：
        模型假设最后一个时间步的表示（context）包含了整个历史序列的关键信息，适用于短期预测。若需捕获长期依赖，可考虑使用所有时间步的输出（如通过注意力机制聚合）。
        
        多步预测机制：
        通过单个上下文向量预测未来多个时间步，这种设计简化了模型结构，但可能限制了对复杂时序关系的建模能力。
        '''

        # 通过解码器生成预测
        output = self.decoder(context)
        
        # 重塑输出为 [batch_size, pred_length, output_dim]
        output = output.view(batch_size, self.pred_length, -1)
        
        return output



def evaluate_model(model, data_loader, criterion, device):
    """评估模型性能
    model：待评估的模型（如之前定义的TransformerModel）。
    data_loader：数据加载器（DataLoader实例），提供待评估的数据集（如验证集）。
    criterion：损失函数（与训练时一致，如nn.MSELoss），用于计算预测误差。
    device：计算设备（cuda或cpu），确保数据和模型在同一设备上。
    """
    model.eval() # 将模型切换为评估模式（与训练模式model.train()对应）
    '''
    具体影响：
        关闭 Dropout 层（不再随机丢弃神经元）。
        固定 BatchNorm 层的统计量（使用训练时的均值和方差，而非实时计算）。
        确保评估过程的稳定性和可重复性。
    '''

    total_loss = 0.0
    with torch.no_grad():
        '''
        torch.no_grad() 是一个上下文管理器，进入该上下文后：
            禁用自动求导（不再计算张量的梯度）。
            减少内存占用（无需存储梯度相关信息）。
            加快评估速度（省去反向传播的计算开销）。
        '''

        for x, y, stock_idx in data_loader:
            x, y = x.to(device), y.to(device) # 1. 数据移至目标设备
            y_pred = model(x)  # 2. 模型预测（前向传播）
            loss = criterion(y_pred, y) # 3. 计算当前批次的损失
            total_loss += loss.item() * x.size(0)  # 4. 累计总损失（按样本数量加权）
    return total_loss / len(data_loader.dataset) # 总损失 / 总样本数 = 平均损失

def train_model(model, train_loader, val_loader, criterion, optimizer, epochs, device, scheduler=None,model_path='best_transformer_model.pth'):
    """训练Transformer模型
        model：待训练的 Transformer 模型（TransformerModel实例）。
        train_loader训练集的数据加载器（DataLoader）.
        val_loader验证集的数据加载器（DataLoader）.
        criterion：损失函数（如nn.MSELoss，用于衡量预测误差），输入预测与标签，输出损失。
        optimizer：参数指向 model.parameters()的 优化器（如Adam，用于更新模型参数）。
        epochs：训练轮数（完整遍历训练集的次数）。
        device：训练设备（cuda或cpu）。
        scheduler：优化器指向 optimizer 的学习率调度器（可选，用于每轮训练动态调整学习率）。
            
        例如：
            optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=5)
            使用ReduceLROnPlateau调度器时，若验证损失不再下降，会自动降低学习率（如乘以 0.5），帮助模型跳出局部最优。

        返回(train_losses, val_losses)，元组的元素是损失函数列表，可用于后续绘制损失曲线，分析模型收敛情况。
    """
    model.train()       # 切换模型至训练模式（启用Dropout等）
    train_losses = []   # 记录每轮训练损失
    val_losses = []     # 记录每轮验证损失
    best_val_loss = float('inf') # 最佳验证损失（初始化为无穷大），用于后续保存。
    
    '''
    model.train() 确保模型中的 dropout 层和批归一化层处于训练状态（与推理模式model.eval()对应）。
    '''

    for epoch in range(epochs): # epochs是训练轮数，每轮回遍历所有批次数据。
        train_loss = 0.0
        model.train()  # 确保每轮都处于训练模式
        
        # 使用tqdm显示训练进度条，方便实时观察训练进度。
        # x, y, stock_idx 取自 train_loader ，每个批次都做一次反向传播，跟新参数
        for x, y, stock_idx in tqdm(train_loader, desc=f'Epoch {epoch+1}/{epochs} [Train]'):
            '''
            train_loader：迭代对象（通常是 PyTorch 的 DataLoader），包含训练数据的批次。tqdm 会遍历该对象并跟踪进度。

                desc=f'Epoch {epoch+1}/{epochs} [Train]'：进度条的描述文字，用于区分不同的训练阶段。这里显示当前是第几个 epoch（epoch+1）、总 epoch 数（epochs），以及当前是训练阶段（[Train]）。

                执行时，控制台会显示类似这样的进度条：
                Epoch 1/10 [Train]: 100%|███| 120/120 [00:23<00:00,  5.12it/s]

                100%：完成进度。
                ████：可视化进度条
                120/120：当前已处理 120 个批次，总共有 120 个批次。
                00:23<00:00：已耗时 23 秒，预计剩余时间 0 秒（已完成）。
                5.12it/s：每秒处理 5.12 个批次（迭代速度）。

                tqdm 的工作原理
                tqdm 本质是一个迭代器包装器，它会 “包裹” 原始迭代对象（如 train_loader）。每次从 tqdm 迭代器中获取一个批次数据时，它会自动更新进度条状态（完成百分比、耗时等）。当迭代结束时，进度条会显示最终统计信息（总耗时、平均速度等）。
            '''
            x, y = x.to(device), y.to(device) # 数据移至设备
            optimizer.zero_grad()  # 清空梯度
            y_pred = model(x) # 前向传播
            loss = criterion(y_pred, y)# 计算损失
            loss.backward() # 反向传播
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)# 将梯度范数限制在 1.0 以内，防止梯度爆炸（序列模型常见问题）。
            optimizer.step() #  根据梯度更新模型参数。    
            train_loss += loss.item() * x.size(0) # 按样本数量加权累计损失，便于后续计算平均损失。
        
        train_loss /= len(train_loader.dataset)  # 总损失 / 总样本数
        train_losses.append(train_loss)   # 记录当前轮次的平均损失

        '''
        DataLoader 与 Dataset 的关系:
            DataLoader 本身不存储样本，它只是对 dataset 进行封装，提供批量加载、多进程加速等功能。
            train_loader.dataset 直接指向原始数据集（如 train_dataset）
            len(train_loader.dataset) 等价于 len(train_dataset)，即调用你自定义 Dataset 类的 __len__ 方法，得到总样本数。

        使用示例：
        train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=num_workers)
            num_workers是加载数据使用的线程数量。
            shuffle表示是否打乱数据。
            batch_size (int): 每个小批量包含的样本数量
        '''
        
        # 调用evaluate_model函数（独立实现）在验证集上评估模型，
        val_loss = evaluate_model(model, val_loader, criterion, device)
        val_losses.append(val_loss)
        # 打印训练进度
        print(f'Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}')
        
        if scheduler: # 根据验证损失调整学习率
            scheduler.step(val_loss) 
        
        if val_loss < best_val_loss: #仅保存验证损失最低的模型参数
            best_val_loss = val_loss
            torch.save(model.state_dict(), model_path)
    
    #返回损失列表，可用于后续绘制损失曲线，分析模型收敛情况。
    return train_losses, val_losses 




def predict_future(model, data, seq_length, pred_length, device, scaler_list):
    """预测未来多个时间步的数据
        model：训练好的 Transformer 模型（TransformerModel实例）。
        data：原始输入数据（通常是标准化后的股票数据列表，列表每个元素是一支股票数据，每支股票数据形状为[天数, 特征数]）。
        seq_length：模型输入的序列长度（如过去 30 天的数据）。
        pred_length：需要预测的未来时间步数（如未来 5 天）。
        device：计算设备（cuda或cpu）。
        scaler_list：每支股票对应的标准化器列表（用于将预测结果反标准化为原始数据范围）。
            
    """
    model.eval()  # 切换至评估模式（禁用Dropout等）
    all_predictions = []  # 存储所有股票的预测结果
    
    with torch.no_grad():  # 关闭梯度计算，加速预测
        # 对每支股票进行预测
        for stock_idx, stock_data in enumerate(data):
            '''
            stock_idx：股票索引（用于匹配对应的标准化器）
            stock_data：单支股票的标准化数据，形状为[总天数, 特征数]（如[3650, 6]）
            data 是股票数据列表，里面每个元素都是一支股票序列。
            '''

            # 准备输入数据（最后seq_length天的数据）
            input_data = stock_data[-seq_length:].copy()
            input_tensor = torch.FloatTensor(input_data).unsqueeze(0).to(device)  # 转换为张量并添加批次维度（模型要求输入形状为[batch_size, seq_length, 特征数]） 这里 batch_size = 1
            
            
            predictions = [] # 预测未来pred_length天的数据
            
            for _ in range(pred_length):

                # 1. 模型预测（前向传播）
                pred = model(input_tensor)  # [1, 1, output_dim]

                # 2. 提取预测值（转换为numpy数组并去除冗余维度）
                pred = pred.cpu().numpy()[0, 0]  # 形状：[output_dim]（如预测3个特征，则为[3]）
                
                # 保存预测结果
                predictions.append(pred)
                
                # 4. 更新输入序列：用新预测结果替换最旧的历史数据，保持序列长度不变
                new_input = np.roll(input_data, -1, axis=0)
                '''
                np.roll 是 NumPy 的滚动（循环移位）函数，这里的作用是将序列整体向前移动 1 位，同时将最旧的元素移到末尾。
                input_data： 原始输入序列，形状为 [seq_length, 特征数]
                -1：         滚动步长为 - 1，表示向前滚动 1 位（若为1则向后滚动）。
                axis=0：     沿第 0 维（时间步维度）滚动。
                '''

                # 只更新预测的特征
                new_input[-1, :] = pred   # 最后一行的目标特征替换为新预测值
                # new_input[-1, :]：选取更新后序列的最后一行

                input_data = new_input.copy()  # 更新输入数据
                
                # 更新输入张量 （准备下一次预测）
                input_tensor = torch.FloatTensor(input_data).unsqueeze(0).to(device)
            
            # 将预测结果转换为numpy数组
            predictions = np.array(predictions)
            
            # 反标准化预测结果 
            dim = len(stock_data[0])
            # dummy_array = np.zeros((len(predictions), 6))  # 创建虚拟数组用于反标准化
            dummy_array = np.zeros((len(predictions), dim))  # 创建虚拟数组用于反标准化
            
            # 填充预测的特征（仅替换目标特征列，其他列用0占位）
            dummy_array[:, :] = predictions  # 填充预测的特征

            # 反标准化：使用该股票对应的scaler（恢复为原始数据范围）
            predictions = scaler_list[stock_idx].inverse_transform(dummy_array)[:, :]  # 反标准化

            '''
            模型预测的是标准化后的数据（均值 0、标准差 1），需通过训练时的scaler转换回原始量级.

            '''
            
            all_predictions.append(predictions)
    
    return all_predictions




def plot_predictions(actual_data, predictions, stock_names, feature_names, past_days=30):
    """可视化预测结果"""
    stock_num = len(actual_data)
    feature_num = len(feature_names)
    
    plt.figure(figsize=(15, 4 * stock_num))
    
    for i in range(stock_num):
        # 获取实际数据的最后past_days天
        actual_past = actual_data[i][-past_days:, :]  # 最高价、最低价、成交量
        
        for j in range(feature_num):
            plt.subplot(stock_num, feature_num, i * feature_num + j + 1)
            
            # 绘制实际数据
            actual_days = np.arange(past_days)
            plt.plot(actual_days, actual_past[:, j], label='Actual', color='blue')
            
            # 绘制预测数据
            pred_days = np.arange(past_days, past_days + len(predictions[i]))
            plt.plot(pred_days, predictions[i][:, j], label='Predicted', color='red')
            
            plt.title(f'{stock_names[i]} - {feature_names[j]}')
            plt.xlabel('Days')
            plt.ylabel(feature_names[j])
            plt.legend()
            plt.grid(True)
    
    plt.tight_layout()
    plt.show()

def main_rand():
    # 设置参数
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    seq_length = 60  # 输入序列长度：使用前60天的数据
    pred_length = 10  # 预测序列长度：预测未来10天
    batch_size = 64
    epochs = 50
    input_dim = 6  # 输入特征维度：开盘价、最高价、最低价、收盘价、成交量、成交额
    output_dim = input_dim  # 输出特征维度：最高价、最低价、成交量
    d_model = 128  # Transformer模型维度
    nhead = 8  # 注意力头数
    num_layers = 2  # Transformer层数
    dropout = 0.1  # Dropout率
    lr = 0.0005  # 学习率
    
    # 模拟5支股票的10年日K线数据
    # 实际应用中，应替换为真实的股票数据加载
    # 每支股票数据形状：[天数, 特征数(6)]
    np.random.seed(42)
    stock_data = []
    stock_names = ['Stock A', 'Stock B', 'Stock C', 'Stock D', 'Stock E']
    
    for _ in range(5):
        # 生成随机数据
        days = 365 * 10  # 10年数据
        data = np.zeros((days, 6))
        
        # 生成开盘价
        open_price = np.random.uniform(10, 100, days)
        data[:, 0] = open_price
        
        # 生成最高价（高于开盘价）
        high_price = open_price + np.random.uniform(0, 10, days)
        data[:, 1] = high_price
        
        # 生成最低价（低于开盘价）
        low_price = open_price - np.random.uniform(0, 10, days)
        data[:, 2] = low_price
        
        # 生成收盘价（在最高价和最低价之间）
        close_price = np.random.uniform(low_price, high_price)
        data[:, 3] = close_price
        
        # 生成成交量
        volume = np.random.uniform(1e5, 1e7, days)
        data[:, 4] = volume
        
        # 生成成交额
        amount = close_price * volume
        data[:, 5] = amount
        
        # 添加一些趋势和季节性
        trend = np.linspace(0, 50, days)
        seasonality = 10 * np.sin(np.linspace(0, 20*np.pi, days))
        for i in range(4):  # 对价格类特征添加趋势和季节性
            data[:, i] += trend + seasonality
        
        stock_data.append(data)
    
    # 数据标准化
    scaler_list = []
    scaled_data = []
    
    '''
    这段代码的作用是对多支股票的原始数据进行标准化处理，确保不同股票的特征处于相同的量级，从而提高模型的训练稳定性和预测精度。

    价格类特征（如开盘价、收盘价）：通常在几十到几百元。
    成交量 / 成交额：可能在几十万到几亿。
    这种量级差异会导致模型更关注数值大的特征（如成交量），而忽略数值小的特征（如价格波动）。标准化通过将数据缩放至均值为 0、标准差为 1，消除这种影响。
    '''
    for stock in stock_data:
        scaler = StandardScaler()
        scaled = scaler.fit_transform(stock)
        scaler_list.append(scaler)
        scaled_data.append(scaled)
    
    # 划分训练集和验证集
    train_ratio = 0.8
    train_data = []
    val_data = []
    
    for stock in scaled_data:
        train_size = int(len(stock) * train_ratio)
        train_data.append(stock[:train_size])
        val_data.append(stock[train_size:])
    
    # 创建数据集和数据加载器
    train_dataset = MultiStockDataset(train_data, seq_length, pred_length)
    val_dataset = MultiStockDataset(val_data, seq_length, pred_length)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    
    # 初始化模型
    model = TransformerModel(
        input_dim=input_dim,
        d_model=d_model,
        nhead=nhead,
        num_layers=num_layers,
        output_dim=output_dim,
        pred_length=pred_length,
        dropout=dropout
    ).to(device)
    
    # 定义损失函数和优化器
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, 
        mode='min', 
        factor=0.5, 
        patience=5, 
        verbose=True
    )
    
    # 训练模型
    print("开始训练模型...")
    train_losses, val_losses = train_model(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        criterion=criterion,
        optimizer=optimizer,
        epochs=epochs,
        device=device,
        scheduler=scheduler
    )
    
    # 加载最佳模型进行预测
    model.load_state_dict(torch.load('best_transformer_model.pth'))
    
    # 预测未来10天的数据
    print("开始预测未来10天的数据...")
    predictions = predict_future(
        model=model,
        data=scaled_data,
        seq_length=seq_length,
        pred_length=pred_length,
        device=device,
        scaler_list=scaler_list
    )
    
    """
    # 可视化预测结果
    feature_names = ['High Price', 'Low Price', 'Volume']
    plot_predictions(
        actual_data=[scaler.inverse_transform(data) for scaler, data in zip(scaler_list, scaled_data)],
        predictions=predictions,
        stock_names=stock_names,
        feature_names=feature_names
    )
    """
    
    # 打印预测结果
    for i, stock_name in enumerate(stock_names):
        print(f"\n{stock_name} 未来10天预测:")
        print("日期\t\t最高价\t\t最低价\t\t成交量")
        for j in range(pred_length):
            print(f"Day {j+1}\t{predictions[i][j, 0]:.2f}\t\t{predictions[i][j, 1]:.2f}\t\t{int(predictions[i][j, 2]):,}")

def main():
    # 设置参数
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 筛选 100 支股票，近10年数据，
    stock_num = 100

    model_path='stock_transformer_model.pth'
    seq_length = 120   # 输入序列长度：使用前60天的数据
    pred_length = 5   # 预测序列长度：预测未来10天
    batch_size = 160
    epochs = 20

    input_dim = 7  # 输入特征维度：open    high     low   close    volume(手)          amount(金额)      turn(换手率) 
    output_dim = input_dim  

    d_model = 64  # Transformer模型维度
    nhead = 16  # 注意力头数
    num_layers = 8  # Transformer层数
    dropout = 0.1  # Dropout率
    lr = 0.0005  # 学习率


    stk_data = stock_loader.stock_data()
    stk_data.directory_path = '/root/deep-learning---second-stage/stock_data/'
    #stk_list = stk_data.load_stock_list()

    # 上证50股票列表（包含代码和名称）
    sse50_stocks = [
        ["sh.601127", "赛力斯"], ["sh.601766", "中国中车"], ["sh.601816", "京沪高铁"],
        ["sh.603993", "洛阳钼业"], ["sh.688256", "寒武纪"], ["sh.601288", "农业银行"],
        ["sh.600048", "保利发展"], ["sh.601166", "兴业银行"], ["sh.600030", "中信证券"],
        ["sh.601857", "中国石油"], ["sh.601668", "中国建筑"], ["sh.601628", "中国人寿"],
        ["sh.601398", "工商银行"], ["sh.601318", "中国平安"], ["sh.601088", "中国神华"],
        ["sh.600519", "贵州茅台"], ["sh.600036", "招商银行"], ["sh.600028", "中国石化"],
        ["sh.600887", "伊利股份"], ["sh.600309", "万华化学"], ["sh.600276", "恒瑞医药"],
        ["sh.600690", "海尔智家"], ["sh.601888", "中国中免"], ["sh.603259", "药明康德"],
        ["sh.600031", "三一重工"], ["sh.601012", "隆基绿能"], ["sh.603501", "韦尔股份"],
        ["sh.600809", "山西汾酒"], ["sh.601899", "紫金矿业"], ["sh.600900", "长江电力"],
        ["sh.601919", "中远海控"], ["sh.600406", "国电南瑞"], ["sh.601225", "陕西煤业"],
        ["sh.600050", "中国联通"], ["sh.601728", "中国电信"], ["sh.688111", "金山办公"],
        ["sh.688041", "海光信息"], ["sh.688981", "中芯国际"], ["sh.601988", "中国银行"],
        ["sh.601601", "中国太保"], ["sh.600150", "中国船舶"], ["sh.601985", "中国核电"],
        ["sh.688012", "中微公司"], ["sh.601658", "邮储银行"], ["sh.601328", "交通银行"],
        ["sh.600941", "中国移动"], ["sh.601211", "国泰海通"], ["sh.688008", "澜起科技"],
        ["sh.601600", "中国铝业"], ["sh.600760", "中航沈飞"]
    ]

    # 深证50股票列表（包含代码和名称）
    szse50_stocks = [
        ["sz.002180", "纳思达"], ["sz.002252", "上海莱士"], ["sz.002463", "沪电股份"],
        ["sz.300502", "新易盛"], ["sz.002475", "立讯精密"], ["sz.002555", "三七互娱"],
        ["sz.300059", "东方财富"], ["sz.002594", "比亚迪"], ["sz.300015", "爱尔眼科"],
        ["sz.300896", "爱美客"], ["sz.300760", "迈瑞医疗"], ["sz.300750", "宁德时代"],
        ["sz.300498", "温氏股份"], ["sz.300274", "阳光电源"], ["sz.300124", "汇川技术"],
        ["sz.000002", "万科A"], ["sz.000063", "中兴通讯"], ["sz.000100", "TCL科技"],
        ["sz.000333", "美的集团"], ["sz.000538", "云南白药"], ["sz.000568", "泸州老窖"],
        ["sz.000651", "格力电器"], ["sz.000725", "京东方A"], ["sz.000792", "盐湖股份"],
        ["sz.000858", "五粮液"], ["sz.000938", "紫光股份"], ["sz.000977", "浪潮信息"],
        ["sz.001979", "招商蛇口"], ["sz.002001", "新和成"], ["sz.002027", "分众传媒"],
        ["sz.002230", "科大讯飞"], ["sz.002271", "东方雨虹"], ["sz.002311", "海大集团"],
        ["sz.002352", "顺丰控股"], ["sz.000625", "长安汽车"], ["sz.002371", "北方华创"],
        ["sz.002415", "海康威视"], ["sz.002714", "牧原股份"], ["sz.300308", "中际旭创"],
        ["sz.300979", "华利集团"], ["sz.000786", "北新建材"], ["sz.301236", "软通动力"],
        ["sz.300418", "昆仑万维"], ["sz.002422", "科伦药业"], ["sz.000895", "双汇发展"],
        ["sz.300408", "三环集团"], ["sz.000661", "长春高新"], ["sz.002648", "卫星化学"],
        ["sz.300033", "同花顺"], ["sz.000423", "东阿阿胶"]
    ]
    my_codes = [
        ['sh.600028',''],
        ['sh.600089',''],
        ['sh.600348',''],
        ['sh.600694',''],
        ['sh.600985',''],
        ['sh.601868',''],
        ['sz.000063',''],
        ['sz.000591','']]    
    
    all_stocks = sse50_stocks+szse50_stocks+my_codes


    ls = stk_data.load_stock_data(all_stocks,'d')
    #  (股票代码, dataframe) 的列表
    print('load data finish')
    # return

    stock_data = []
    #stock_names = [stock[1] for stock in all_stocks] 
    stock_names = [] 

    for (code ,df) in ls:
        stock_names.append(code)
        stock_data.append(df.apply(pd.to_numeric, errors='coerce').to_numpy()[ : ,1 : 8] )

        # df.apply(pd.to_numeric, errors='coerce') 将df字符串形式的数据改成数值型，errors='coerce'表示不能转换的用nan替换。
        # .to_numpy() 将df 转成numpy， [ : ,1 : 6] 去掉第一列（第一列是日期，无法转换成数字，numpy中是nan）

    # 数据标准化
    scaler_list = []
    scaled_data = []
    '''
    这段代码的作用是对多支股票的原始数据进行标准化处理，确保不同股票的特征处于相同的量级，从而提高模型的训练稳定性和预测精度。

    价格类特征（如开盘价、收盘价）：通常在几十到几百元。
    成交量 / 成交额：可能在几十万到几亿。
    这种量级差异会导致模型更关注数值大的特征（如成交量），而忽略数值小的特征（如价格波动）。标准化通过将数据缩放至均值为 0、标准差为 1，消除这种影响。
    '''
    for stock in stock_data:
        scaler = StandardScaler()
        scaled = scaler.fit_transform(stock)
        scaler_list.append(scaler)
        scaled_data.append(scaled)
    

    if 1 :  # 训练

        # 划分训练集和验证集
        train_ratio = 0.8  # 训练集比例。序列前面百分之 80% 作为u训练集
        train_data = []
        val_data = []
        
        for stock in scaled_data:
            train_size = int(len(stock) * train_ratio)
            train_data.append(stock[:train_size])
            val_data.append(stock[train_size:])
        
        # 创建数据集和数据加载器
        train_dataset = MultiStockDataset(train_data, seq_length, pred_length)
        val_dataset = MultiStockDataset(val_data, seq_length, pred_length)
        
        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=batch_size)
        
        # 初始化模型
        model = TransformerModel(
            input_dim=input_dim,
            d_model=d_model,
            nhead=nhead,
            num_layers=num_layers,
            output_dim=output_dim,
            pred_length=pred_length,
            dropout=dropout
        ).to(device)
        
        # 定义损失函数和优化器
        criterion = nn.MSELoss()
        optimizer = optim.Adam(model.parameters(), lr=lr)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, 
            mode='min', 
            factor=0.5, 
            patience=5, 
            verbose=True
        )
        # 训练模型
        print("开始训练模型...")
        train_losses, val_losses = train_model(
            model=model,
            train_loader=train_loader,
            val_loader=val_loader,
            criterion=criterion,
            optimizer=optimizer,
            epochs=epochs,
            device=device,
            scheduler=scheduler,
            model_path=model_path
        )
        
    if 0 : # 预测

        # 加载最佳模型进行预测
        model.load_state_dict(torch.load(model_path))
        
        # 预测未来10天的数据
        print("开始预测未来10天的数据...")
        predictions = predict_future(
            model=model,
            data=scaled_data,
            seq_length=seq_length,
            pred_length=pred_length,
            device=device,
            scaler_list=scaler_list
        )
        
        # 打印预测结果
        for i, stock_name in enumerate(stock_names):
            print(f"\n{stock_name} 未来10天预测:")
            print("日期\t\t最高价\t\t最低价\t\t成交量")
            for j in range(pred_length):
                print(f"Day {j+1}\t{predictions[i][j, 0]:.2f}\t\t{predictions[i][j, 1]:.2f}\t\t{int(predictions[i][j, 2]):,}")


if __name__ == "__main__":
    main()




